#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
客户端端口复用 2.0  ––  仅出站方向改写
python3 reuse_v2.py
"""
import json, time, socket, struct, psutil, subprocess, threading
import pyshark
from scapy.all import IP, TCP, UDP, raw
import netfilterqueue as nfq

CONF_FILE = "config.json"
FIRST_ONLY = True          # True=仅第一个包追加2-byte；False=每个包都追加
_added_port = set()        # 已追加过2-byte的本地端口集合

# ---------- 1. 读配置 ----------
with open(CONF_FILE, encoding="utf-8") as f:
    cfg = json.load(f)
SERVER_IP = cfg["server_ip"]
SERVER_A_PORT = cfg["A_port"]
SERVER_B_PORT = cfg["B_port"]

def get_target_pid():
    if cfg.get("pid"):
        return cfg["pid"]
    for p in psutil.process_iter(['pid', 'name']):
        if p.info['name'] == cfg["name"]:
            return p.info['pid']
    raise RuntimeError("进程未找到")

PID = get_target_pid()
print(f"[+] 目标进程 PID = {PID}")

# ---------- 2. 拿到该进程当前所有本地 TCP 端口 ----------
def get_ports_by_pid(pid):
    ports = set()
    for c in psutil.Process(pid).connections():
        if c.status == "ESTABLISHED" and c.laddr:
            ports.add(c.laddr.port)
    return sorted(ports)

LOCAL_PORTS = get_ports_by_pid(PID)
print(f"[+] 进程正在使用的本地端口: {LOCAL_PORTS}")

# ---------- 3. 构造 pyshark 抓包过滤器 ----------
# 构造BPF过滤器字符串，仅抓取目标进程本地端口的出站TCP/UDP数据包
bpf_filter = " or ".join([f"(tcp and src port {p}) or (udp and src port {p})"
                          for p in LOCAL_PORTS])

# ---------- 4. 用 netfilterqueue 在 OUTPUT 链做改写 ----------
def handle_packet(nf_pkt):
    """
    nf_pkt 是 netfilterqueue 对象，payload 为原始 IP 报文
    """
    # 使用Scapy解析IP层
    ip = IP(nf_pkt.get_payload())
    # 若源地址是服务器IP，不做处理直接放行（防止回环）
    if ip.src == SERVER_IP:
        nf_pkt.accept()
        return

    L3 = ip
    L4 = L3.payload
    sport = None
    # 判断传输层协议并提取源端口和目的端口
    if TCP in L3:
        sport = L3[TCP].sport
        dport = L3[TCP].dport
    elif UDP in L3:
        sport = L3[UDP].sport
        dport = L3[UDP].dport
    else:
        # 非TCP/UDP包直接放行
        nf_pkt.accept()
        return

    # 只处理我们关心的本地端口
    if sport not in LOCAL_PORTS:
        nf_pkt.accept()
        return

    # 是否已追加过 2-byte
    need_append = (FIRST_ONLY and sport not in _added_port) or (not FIRST_ONLY)
    if need_append:
        _added_port.add(sport)
        # 在负载尾部追加真实目的端口（2-byte 小端）
        tail = struct.pack("<H", dport)
        if TCP in L3:
            # 将TCP负载与追加数据合并
            L3[TCP].payload = raw(L3[TCP].payload) + tail
            # 删除原有校验和，让Scapy自动重新计算
            del L3[TCP].chksum
            del L3.chksum
        else:
            # 将UDP负载与追加数据合并
            L3[UDP].payload = raw(L3[UDP].payload) + tail
            # 删除原有校验和，让Scapy自动重新计算
            del L3[UDP].chksum
            del L3.chksum

    # 统一把目的端口改写成服务器 A 端口
    if TCP in L3:
        L3[TCP].dport = SERVER_A_PORT
    else:
        L3[UDP].dport = SERVER_A_PORT

    # 删除IP层长度字段，让Scapy自动重新计算
    del L3.len
    # 将修改后的报文设置回netfilterqueue并放行
    nf_pkt.set_payload(bytes(L3))
    nf_pkt.accept()

# ---------- 5. 启动 nfqueue ----------
def nfqueue_thread():
    q = nfq.NetfilterQueue()
    q.bind(0, handle_packet)
    print("[+] nfqueue 已绑定，开始接管 OUTPUT 链...")
    q.run()

# 清空mangle表OUTPUT链已有规则
subprocess.call("iptables -t mangle -F OUTPUT", shell=True)
# 对本地端口列表的TCP包设置NFQUEUE
subprocess.call("iptables -t mangle -A OUTPUT -p tcp -m multiport --sports "
                f"{','.join(map(str, LOCAL_PORTS))} -j NFQUEUE --queue-num 0", shell=True)
# 对本地端口列表的UDP包设置NFQUEUE
subprocess.call("iptables -t mangle -A OUTPUT -p udp -m multiport --sports "
                f"{','.join(map(str, LOCAL_PORTS))} -j NFQUEUE --queue-num 0", shell=True)

# 启动nfqueue线程
threading.Thread(target=nfqueue_thread, daemon=True).start()

# ---------- 6. 可选：pyshark 实时打印，仅做可视化 ----------
def shark_thread():
    # 创建pyshark实时抓包对象，使用之前构造的BPF过滤器
    cap = pyshark.LiveCapture(interface=None, bpf_filter=bpf_filter)
    print("[+] pyshark 开始抓包...")
    for pkt in cap.sniff_continuously():
        try:
            # 打印抓到的数据包源IP:端口 -> 目的IP:端口
            print(f"[shark] {pkt.ip.src}:{pkt[pkt.transport_layer].srcport} -> "
                  f"{pkt.ip.dst}:{pkt[pkt.transport_layer].dstport}")
        except Exception as e:
            # 解析异常时忽略
            pass

# 启动抓包线程
threading.Thread(target=shark_thread, daemon=True).start()

# ---------- 7. 主循环 ----------
print("[+] 端口复用 2.0 正在运行，Ctrl-C 退出")
try:
    while True:
        time.sleep(1)
except KeyboardInterrupt:
    print("\n[-] 清理 iptables 并退出")
    # 清空mangle表OUTPUT链规则
    subprocess.call("iptables -t mangle -F OUTPUT", shell=True)